# -*- coding: utf-8 -*-
import numpy as np
import pickle as pkl
from tqdm import tqdm
from datetime import timedelta
from torch.utils.data import Dataset, DataLoader
import torch.nn as nn
import time
import torch
from sklearn import metrics
from sklearn.model_selection import train_test_split
import seaborn as sns
from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay
import matplotlib.pyplot as plt
from sklearn.metrics import precision_score, accuracy_score, f1_score, recall_score


# 超参数设置
data_path =  './data/data.txt'              # 数据集
vocab_path = './data/vocab.pkl'             # 词表
save_path = './saved_dict/lstm.ckpt'        # 模型训练结果
embedding_pretrained = \
    torch.tensor(
    np.load(
        './data/embedding_Tencent.npz')
    ["embeddings"].astype('float32'))
                                            # 预训练词向量
embed = embedding_pretrained.size(1)        # 词向量维度
dropout = 0.5                               # 随机丢弃
num_classes = 2                             # 类别数
num_epochs = 30                             # epoch数
batch_size = 128                            # mini-batch大小
pad_size = 50                               # 每句话处理成的长度(短填长切)
learning_rate = 1e-3                        # 学习率
hidden_size = 128                           # lstm隐藏层
num_layers = 2                              # lstm层数
MAX_VOCAB_SIZE = 10000                      # 词表长度限制
UNK, PAD = '<UNK>', '<PAD>'                 # 未知字，padding符号


def get_data():
    tokenizer = lambda x: [y for y in x]  #  这里定义了一个lambda函数，用于将输入的文本按字符进行分词，将每个字符作为一个token。
    vocab = pkl.load(open(vocab_path, 'rb'))#这行代码加载了一个预先保存的词汇表（vocabulary），通过读取二进制文件的方式加载，其中 vocab_path 是词汇表文件的路径。
    # print('tokenizer',tokenizer)
    print('vocab',vocab)#这行代码打印了加载的词汇表，即包含了哪些字符和它们对应的ID。
    print(f"Vocab size: {len(vocab)}")#这行代码打印了词汇表的大小，即词汇表中包含多少个不同的字符。

    train,dev,test = load_dataset(data_path, pad_size, tokenizer, vocab)
    ''' 这行代码调用了 load_dataset() 函数，传入了数据集路径
     data_path、填充大小 pad_size、分词器 tokenizer 和词汇表 vocab，
     并将返回的训练集、验证集和测试集分配给了 train、dev 和 test 变量。'''
    return vocab, train, dev, test#最后，函数返回了词汇表 vocab 和分割好的训练集、验证集和测试集。

#函数将数据集分割为训练集、验证集和测试集
def load_dataset(path, pad_size, tokenizer, vocab):
    '''
    将路径文本文件分词并转为三元组返回
    :param path: 文件路径
    :param pad_size:  每个序列的大小，用于将序列填充或截断到相同的长度。
    :param tokenizer: 分词器，一个函数，用于将文本分割成tokens。在这里，tokenizer 参数传入的是一个函数对象，用于将文本内容分割成单个字符。
    :param vocab: 词汇表，一个字典，将每个字符映射到对应的ID。
    :return: 二元组，含有字ID，标签
    '''
    contents = []
    n=0
    with open(path, 'r', encoding='gbk') as f:
        # tqdm可以看进度条
        for line in tqdm(f):
            # 默认删除字符串line中的空格、’\n’、't’等。
            lin = line.strip()
            if not lin:
                continue
            # print(lin)
            label,content = lin.split('	####	')
            # word_line存储每个字的id
            words_line = []
            # 分割器，分词每个字
            token = tokenizer(content)
            # print(token)
            # 字的长度
            seq_len = len(token)
            if pad_size:
                # 如果字长度小于指定长度，则填充，否则截断
                if len(token) < pad_size:
                    token.extend([vocab.get(PAD)] * (pad_size - len(token)))
                else:
                    token = token[:pad_size]
                    seq_len = pad_size
            # 将每个字映射为ID
            # 如果在词表vocab中有word这个单词，那么就取出它的id；
            # 如果没有，就去除UNK（未知词）对应的id，其中UNK表示所有的未知词（out of vocab）都对应该id
            for word in token:
                words_line.append(vocab.get(word, vocab.get(UNK)))
            n+=1
            contents.append((words_line, int(label)))

    train, X_t = train_test_split(contents, test_size=0.4, random_state=42)
    dev,test= train_test_split(X_t, test_size=0.5, random_state=42)
    return train,dev,test
# get_data()

class TextDataset(Dataset):
    def __init__(self, data):
        '''类的初始化函数，接受一个参数 data，即包含文本和标签的数据。'''
        self.device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
        self.x = torch.LongTensor([x[0] for x in data]).to(self.device)# 将数据中的文本部分转换为LongTensor类型，并将其移动到之前检查到的设备（GPU或CPU）上。
        self.y = torch.LongTensor([x[1] for x in data]).to(self.device)# 将数据中的标签部分转换为LongTensor类型，并将其移动到之前检查到的设备上。
    def __getitem__(self,index):#定义了 __getitem__ 方法，用于获取数据集中指定索引的数据。
        self.text = self.x[index]#获取指定索引处的文本数据。
        self.label = self.y[index]#获取指定索引处的标签数据。
        return self.text, self.label# 返回文本数据和标签数据的元组。
    def __len__(self):
        return len(self.x)#用于返回数据集的长度（即数据样本的数量）。

# 以上是数据预处理的部分

def get_time_dif(start_time):
    """获取已使用时间"""
    end_time = time.time()
    time_dif = end_time - start_time
    return timedelta(seconds=int(round(time_dif)))


# 定义LSTM模型
class Model(nn.Module):
    def __init__(self):#初始化函数，在创建 Model 类的实例时调用。
        super(Model, self).__init__()
        # 使用预训练的词向量模型，freeze=False 表示允许参数在训练中更新
        # 在NLP任务中，当我们搭建网络时，第一层往往是嵌入层，对于嵌入层有两种方式初始化embedding向量，
        # 一种是直接随机初始化，另一种是使用预训练好的词向量初始化。
        self.embedding = nn.Embedding.from_pretrained(embedding_pretrained, freeze=False)
        # bidirectional=True表示使用的是双向LSTM
        '''
         创建一个嵌入层（Embedding Layer），用于将输入的token序列映射为词向量表示。
         在这里，使用了预训练的词向量模型 embedding_pretrained，并设置 freeze=False 表示允许在训练中更新参数，即不固定词向量的参数。'''
        self.lstm = nn.LSTM(embed, hidden_size, num_layers,
                            bidirectional=True, batch_first=True, dropout=dropout)
        # 因为是双向LSTM，所以层数为config.hidden_size * 2
        '''
        创建一个双向LSTM（Bidirectional LSTM）层。其中 embed 是词向量的维度，hidden_size 是LSTM隐藏层的大小，num_layers
         是LSTM的层数，bidirectional=True 表示使用双向LSTM，batch_first=True 表示输入数据的第一个维度是batch size，dropout=dropout 表示设置dropout的概率。'''
        self.fc = nn.Linear(hidden_size * 2, num_classes)
        '''创建一个全连接（Fully Connected）层，用于将LSTM的输出映射到分类标签空间。输入大小为 hidden_size * 2，
        即双向LSTM的输出大小乘以2，输出大小为 num_classes，即分类的类别数量。'''

    def forward(self, x):#前向传播函数，定义了模型的数据流向。
        out = self.embedding(x)
        # 将输入的token序列 x 通过嵌入层 self.embedding 进行词嵌入，得到词向量表示 out。
        out, _ = self.lstm(out)
        #将词向量表示 out 输入到LSTM层 self.lstm 中进行序列建模，得到输出 out。此处用 _ 来接收LSTM的状态，但在这个任务中并没有使用到这个状态。
        out = self.fc(out[:, -1, :])
        #将LSTM的输出取出序列中最后一个时间步的隐藏状态，然后通过全连接层 self.fc 进行分类预测，得到最终的输出结果 out。
        return out
def get_time_dif(start_time):
    """获取已使用时间"""
    end_time = time.time()
    time_dif = end_time - start_time
    return timedelta(seconds=int(round(time_dif)))

# 权重初始化，默认xavier
# xavier和kaiming是两种初始化参数的方法
def init_network(model, method='xavier', exclude='embedding'):
    #用于初始化模型的参数，可以选择不同的初始化方法，并且可以排除某些参数不进行初始化。
    '''model: 是要初始化参数的模型对象。

method='xavier': 是选择参数初始化方法的参数，默认为 Xavier 初始化方法。

exclude='embedding': 是要排除的参数名称，即不对该名称中包含的参数进行初始化，默认为嵌入层（embedding）参数。'''
    for name, w in model.named_parameters():
        if exclude not in name:
            if 'weight' in name:
                if method == 'xavier':
                    nn.init.xavier_normal_(w)
                elif method == 'kaiming':
                    nn.init.kaiming_normal_(w)
                else:
                    nn.init.normal_(w)
            elif 'bias' in name:
                nn.init.constant_(w, 0)
            else:
                pass

def plot_acc(train_acc):# 是一个列表，包含了每个epoch的训练准确率。
    sns.set(style='darkgrid')#设置绘图风格为暗色网格风格。
    plt.figure(figsize=(10, 7))#创建一个图像对象，设置图像大小为10x7。
    x = list(range(len(train_acc)))#生成一个从0到epoch数的整数列表，作为横坐标。
    plt.plot(x, train_acc, alpha=0.9, linewidth=2, label='train acc')# 绘制训练准确率随epoch变化的曲线，设置线条透明度为0.9，线宽为2，标签为'train acc'。
    plt.xlabel('Epoch')# 设置横坐标轴的标签为'Epoch'。
    plt.ylabel('Acc')# 设置纵坐标轴的标签为'Acc'。
    plt.legend(loc='best')#添加图例，并设置图例位置为最佳位置。
    plt.savefig('results/acc.png', dpi=400)

def plot_loss(train_loss):
    sns.set(style='darkgrid')
    plt.figure(figsize=(10, 7))
    x = list(range(len(train_loss)))
    plt.plot(x, train_loss, alpha=0.9, linewidth=2, label='train acc')
    plt.xlabel('Epoch')
    plt.ylabel('loss')
    plt.legend(loc='best')
    plt.savefig('results/loss.png', dpi=400)

# 定义训练的过程
def train( model, dataloaders):
    '''
    训练模型
    :param model: 模型
    :param dataloaders: 处理后的数据，包含trian,dev,test
    '''
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
    loss_function = torch.nn.CrossEntropyLoss()

    dev_best_loss = float('inf')

    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
    print("Start Training...\n")
    plot_train_acc = []
    plot_train_loss = []

    for i in range(num_epochs):
        # 1，训练循环----------------------------------------------------------------
        # 将数据全部取完
        # 记录每一个batch
        step = 0
        train_lossi=0
        train_acci = 0
        for inputs, labels in dataloaders['train']:
            # 训练模式，可以更新参数
            model.train()
            # print(inputs.shape)
            inputs = inputs.to(device)
            labels = labels.to(device)
            # 梯度清零，防止累加
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = loss_function(outputs, labels)
            loss.backward()
            optimizer.step()

            step += 1
            true = labels.data.cpu()
            predic = torch.max(outputs.data, 1)[1].cpu()
            train_lossi += loss.item()
            train_acci += metrics.accuracy_score(true, predic)
            # 2，验证集验证----------------------------------------------------------------
        dev_acc, dev_loss = dev_eval(model, dataloaders['dev'], loss_function,Result_test=False)
        if dev_loss < dev_best_loss:
            dev_best_loss = dev_loss
            torch.save(model.state_dict(), save_path)
        train_acc = train_acci/step
        train_loss = train_lossi/step
        plot_train_acc.append(train_acc)
        plot_train_loss.append(train_loss)
        print("epoch = {} :  train_loss = {:.3f}, train_acc = {:.2%}, dev_loss = {:.3f}, dev_acc = {:.2%}".
                  format(i+1, train_loss, train_acc, dev_loss, dev_acc))
    plot_acc(plot_train_acc)
    plot_loss(plot_train_loss)
    # 3，验证循环----------------------------------------------------------------
    model.load_state_dict(torch.load(save_path))
    model.eval()
    start_time = time.time()
    test_acc, test_loss = dev_eval(model, dataloaders['test'], loss_function,Result_test=True)
    print('================'*8)
    print('test_loss: {:.3f}      test_acc: {:.2%}'.format(test_loss, test_acc))
'''
初始化优化器和损失函数：使用 Adam 优化器和交叉熵损失函数。
初始化变量 dev_best_loss 为正无穷，用于记录验证集上的最佳损失。
根据是否有 GPU，确定设备类型。
迭代训练 num_epochs 次：
在每个epoch内，迭代训练数据加载器中的每个 batch：
将模型置为训练模式，清零梯度，将输入和标签转移到设备上。
将输入传入模型，计算输出，计算损失并进行反向传播优化。
记录训练损失和准确率。
在每个epoch结束后，使用验证集评估模型性能，计算验证损失和准确率。
如果验证损失低于 dev_best_loss，则更新 dev_best_loss 并保存模型参数。
记录训练过程中的损失和准确率，并打印当前 epoch 的训练和验证结果。
调用 plot_acc() 和 plot_loss() 函数绘制训练过程中的准确率和损失曲线。
加载验证损失最佳的模型参数，并评估测试集上的性能，打印测试结果。'''
def result_test(real, pred):
    cv_conf = confusion_matrix(real, pred)
    acc = accuracy_score(real, pred)
    precision = precision_score(real, pred, average='micro')
    recall = recall_score(real, pred, average='micro')
    f1 = f1_score(real, pred, average='micro')
    patten = 'test:  acc: %.4f   precision: %.4f   recall: %.4f   f1: %.4f'
    print(patten % (acc, precision, recall, f1,))
    labels11 = ['negative', 'active']
    disp = ConfusionMatrixDisplay(confusion_matrix=cv_conf, display_labels=labels11)
    disp.plot(cmap="Blues", values_format='')
    plt.savefig("results/reConfusionMatrix.tif", dpi=400)
'''
将模型置为评估模式，这样模型的行为就不会被自动求导所记录。
初始化 loss_total 为 0，用于累计总损失。
初始化 predict_all 和 labels_all 为空数组，用于保存所有预测结果和真实标签。
使用 torch.no_grad() 上下文管理器，避免在评估过程中记录梯度信息，以节省内存。
遍历数据集中的每个样本：
将文本输入模型，获取模型的预测输出。
计算当前样本的损失，并累加到 loss_total 中。
将标签转换为 NumPy 数组，并将模型预测的标签转换为 NumPy 数组。
将真实标签和预测标签添加到 labels_all 和 predict_all 中。
计算整个数据集上的准确率。
如果 Result_test 为 True，则调用 result_test() 函数，输出测试结果。
返回准确率和平均损失。'''
# 模型评估
def dev_eval(model, data, loss_function,Result_test=False):
    '''
    :param model: 模型
    :param data: 验证集集或者测试集的数据
    :param loss_function: 损失函数
    :return: 损失和准确率
    '''
    model.eval()
    loss_total = 0
    predict_all = np.array([], dtype=int)
    labels_all = np.array([], dtype=int)
    with torch.no_grad():
        for texts, labels in data:
            outputs = model(texts)
            loss = loss_function(outputs, labels)
            loss_total += loss.item()
            labels = labels.data.cpu().numpy()
            predic = torch.max(outputs.data, 1)[1].cpu().numpy()
            labels_all = np.append(labels_all, labels)
            predict_all = np.append(predict_all, predic)

    acc = metrics.accuracy_score(labels_all, predict_all)
    if Result_test:
        result_test(labels_all, predict_all)
    else:
        pass
    return acc, loss_total / len(data)

if __name__ == '__main__':
    # 设置随机数种子，保证每次运行结果一致，不至于不能复现模型
    np.random.seed(1)#设置随机数种子：通过设置随机数种子，可以确保每次运行的结果一致，这对于模型训练的可重复性很重要。
    torch.manual_seed(1)
    torch.cuda.manual_seed_all(1)
    torch.backends.cudnn.deterministic = True  # 保证每次结果一样

    start_time = time.time()
    print("Loading data...")
    vocab, train_data, dev_data, test_data = get_data()
    dataloaders = {
        'train': DataLoader(TextDataset(train_data), batch_size, shuffle=True),
        'dev': DataLoader(TextDataset(dev_data), batch_size, shuffle=True),
        'test': DataLoader(TextDataset(test_data), batch_size, shuffle=True)
    }
    time_dif = get_time_dif(start_time)
    print("Time usage:", time_dif)
    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
    model = Model().to(device)
    init_network(model)
    train(model, dataloaders)
